from fastmcp import FastMCP  
import asyncio  
from fastmcp import Client  
import akshare as ak  
  
mcp = FastMCP('stock.mcp')  
  
@mcp.tool()  
def get_stock_data(stock_code: str) -> dict:  
    """  
    获取A股个股的实时行情和K线数据  
      
    :param stock_code: 股票代码,例如 "000001" 或 "600000"  
    :result: 返回包含基本信息和K线数据的字典,格式为 {'basic_info': {...}, 'k_data': [...]}  
    """  
    try:  
        # 直接获取K线数据(包含价格信息),避免调用重量级的全市场API  
        k_data = ak.stock_zh_a_hist(symbol=stock_code, period="daily", adjust="qfq")  
          
        if k_data.empty:  
            return {'error': f'未找到股票代码 {stock_code} 的数据'}  
          
        # 从K线数据提取最新行情  
        latest = k_data.iloc[-1]  
        basic_info = {  
            '代码': stock_code,  
            '最新价': float(latest['收盘']),  
            '涨跌幅': round(((latest['收盘'] - latest['开盘']) / latest['开盘'] * 100), 2) if latest['开盘'] > 0 else 0,  
            '成交量': int(latest['成交量']),  
            '日期': str(latest['日期'])  
        }  
          
        return {  
            'basic_info': basic_info,  
            'k_data': k_data.tail(100).to_dict('records')  # 只返回最近100条  
        }  
    except Exception as e:  
        # 添加fallback静态数据,参考index_service.py的做法  
        print(f"获取股票数据失败: {e}")  
        return {  
            'basic_info': {  
                '代码': stock_code,  
                '名称': f'股票{stock_code}',  
                '最新价': 10.50,  
                '涨跌幅': 1.25,  
                '提示': '当前为示例数据,实时数据获取失败'  
            },  
            'k_data': []  
        }  


STOCK_SCHEMA = {  
    "type": "function",  
    "function": {  
        "name": "get_stock_data",  
        "description": "获取A股个股的实时行情和K线数据",  
        "parameters": {  
            "type": "object",  
            "properties": {  
                "stock_code": {  
                    "type": "string",  
                    "description": "股票代码,例如 \"000001\" 或 \"600000\""  
                }  
            },  
            "required": ["stock_code"]  
        }  
    }  
}

async def main():  
    client = Client(mcp)  
    async with client:  
        tools = await client.list_tools()  
        print('可用工具:', tools)  
          
        result = await client.call_tool('get_stock_data', {'stock_code': '000001'})  
        print('调用结果:', result)  
  
if __name__ == '__main__':  
    asyncio.run(main())